package 力扣.堆;

import java.util.*;

public class 最小的k个数Offer40 {
    public static void main(String[] args) {
        int[] ints = {3,2,1};
//        int[] ints = {0,1,2,1};
        int k = 2;
        int[] leastNumbers = getLeastNumbers(ints, k);
        System.out.println(Arrays.toString(leastNumbers));
    }
    public static int[] getLeastNumbers(int[] arr, int k) {
        int[] ans = new int[k];
        PriorityQueue pq = new PriorityQueue<Integer>(new Comparator<Integer>() {//通过重写抽象类Comparator中的compare方法可以更改比较方式
            @Override
            public int compare(Integer o1, Integer o2) {
//                return  o1 - o2;// o1 减去 o2 ：则是小根堆
                return  o2 - o1;//  o2 减去 o1：则是大根堆
            }
        });
        for (int te : arr) {
            if (!pq.isEmpty() && pq.size() >= k){
                if ((int)pq.peek() > te){
                    pq.poll();
                    pq.offer(te);
                }

            }else{
                pq.offer(te);
            }
        }
        int i = 0;
        while (!pq.isEmpty() && i < k){
           ans[i] = (int)pq.poll();
           i++;
        }
        return ans;
    }

    public int[] getLeastNumbers2(int[] arr, int k) {
        if (arr.length == 0 || k == arr.length){
            return new int[0];
        }
        PriorityQueue<Integer> pq = new PriorityQueue<>(new Comparator<Integer>() {
            @Override
            public int compare(Integer o1, Integer o2) {
                return o2 - o1;//大根堆
            }
        });
        for (int te : arr) {
            pq.add(te);
            if (pq.size() > k){
                pq.poll();
            }
        }
        int i = 0;
        int[] ans = new int[k];
        while (!pq.isEmpty()){
            ans[i++] = pq.poll();
        }
        return ans;

    }
}
